# -*- coding: utf-8 -*-
'''
Created on 16.08.2019

@author: yu03
'''


from FFT_Interpolation import *
from mpl_toolkits.mplot3d import Axes3D
from scipy import signal
from scipy.optimize import curve_fit
from scipy.signal import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation


def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
# p= -0.0213125
p= -0.02135
def fit_func_phase(x, q):
    return p*x+q
def fit_func_line(x, p, q):
    return p*x+q

def w_z(w_0, Z_R, Z):
    return w_0*np.sqrt(1 + (Z/Z_R)**2)

def R_z(Z_R, Z):
    return Z*(1+(Z_R/Z)**2)

def Phi_Gouy(Z_R, Z):
    return np.arctan(Z/Z_R)    

j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]
k = 2 * np.pi / Lamda
fs_cam = 1000 ### Hardware Trigger

pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size

img_set = []
hor_phase_centers = []
hor_freq_set = []
hor_f_fit_set = []
hor_phi_fit_set = []
angle_hor_set = []

I_0 = 1023
V_r_x, V_r_y, V_r_z = 0.000, 0.0005, 1 # Reference
V_m_x, V_m_y, V_z_m = 0.001, 0.00, 1 # Measurement
M, N = 0.1, 0.05
# L, D = 1.5, 0
L, D = 0.15, 0
w_0 = 1e-3
Z_R = np.pi * w_0**2 / Lamda
dx = np.linspace((-pix_num/2+1)*pix_size, pix_num/2*pix_size, num=pix_num)
dy = dx
X, Y = np.meshgrid(dx, dy)

###### One frame test
for num in range(600):
    
    Z_p_r = M + N + L * (2/(1-V_r_x**2-V_r_y**2) - 1)
    Z_p_m = M + N + (L+D) * (2/(1-V_m_x**2-V_m_y**2) - 1)
    
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
    
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    
    r_mea = np.sqrt( X**2 + Y**2 + (L+D)**2 - ((2*X*V_m_x+2*Y*V_m_y+(L+D)*(1-V_m_x**2-V_m_y**2))/(1+V_m_x**2+V_m_y**2))**2 )
    r_ref = np.sqrt( X**2 + Y**2 + (L)**2 - ((2*X*V_r_x+2*Y*V_r_y+L*(1-V_r_x**2-V_r_y**2))/(1+V_r_x**2+V_r_y**2))**2 )
    
    R_ref = R_z(Z_R, Z_p_r+d_ref)
    R_mea = R_z(Z_R, Z_p_m+d_mea)
    
    Phi_gouy_ref = Phi_Gouy(Z_R, Z_p_r+d_ref)
    Phi_gouy_mea = Phi_Gouy(Z_R, Z_p_m+d_mea)
    
    diff_phi = k * (diff_Z_p + (d_mea-d_ref) + r_mea**2/2/R_mea - r_ref**2/2/R_ref) + Phi_gouy_mea - Phi_gouy_ref
    A_beat = 0.5 * I_0 * w_0**2 / w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea) * np.exp((-X**2-Y**2)/w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea))
#     A_beat = 0.5 * I_0
    I_beat = A_beat*(1+np.cos(diff_phi))
#     I_beat = I_beat.astype(np.int)
    
    hor_center = I_beat[340+num] ### 340~940
    
    DC_num = 1000
    hor_freq_estim, hor_phase_estim, hor_freqline, hor_sig_magnitude, hor_sig_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase = FFT_interpolation_2(hor_center, pix_size, 1e5, DC_num)
    hor_FFT_start = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][0]+DC_num
    hor_FFT_end = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][-1]+DC_num
    hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
    hor_fit_y = hor_sig_magnitude[hor_FFT_start:hor_FFT_end+1]
    hor_fit_phase = hor_sig_phase[hor_FFT_start:hor_FFT_end+1]
    hor_params = curve_fit(fit_func, hor_fit_x, hor_fit_y)
    [hor_a, hor_b, hor_c] = hor_params[0]
    hor_f_fit = hor_b
    hor_f_fit_set.append(hor_f_fit)
    hor_params = curve_fit(fit_func_phase, hor_fit_x, np.unwrap(hor_fit_phase))
#     [hor_q] = hor_params[0]
#     hor_phi_fit = p*hor_freqline[hor_m_k_num]+hor_q
#     hor_phi_fit_set.append(hor_phi_fit+2*np.pi-0.065)
    # print(num, hor_m_k_num)
       
    hor_freq_set.append(hor_freq_estim)    
    hor_phase_centers.append(hor_phase_estim)
    
    angle_hor = hor_f_fit*Lamda/2 - V_r_x
    angle_hor_set.append(angle_hor)
    print(num, hor_m_k_num, angle_hor, hor_phase_estim)

plt.figure(1)
im = plt.imshow(I_beat, cmap='gray')
plt.colorbar(im, fraction=0.046, pad=0.04)

plt.figure(2)
plt.subplot(2,1,1)
plt.plot(hor_center)
plt.subplot(2,1,2)
plt.plot(hor_freqline, hor_sig_magnitude)

plt.figure(3)
plt.subplot(2,1,1)
plt.plot(angle_hor_set-np.arctan(V_m_x))
plt.ylabel('rad')
plt.subplot(2,1,2)
# plt.plot(hor_phase_centers)
phase_diff = np.diff(np.unwrap(hor_phase_centers))
print(phase_diff)
plt.plot(phase_diff/2/np.pi*Lamda/2*1e9)
plt.ylabel('nm')
plt.show()






